#!/usr/bin/env python

# use this script to create a Makefile to start from RAW files and end
# up with collapsed omssa CSV files

import os, sys, getopt
from xlh import External

def which(filename):
  for dir in os.getenv('PATH').split(':'):
    FilePath = os.path.join(dir, filename)
    if os.path.exists(FilePath):
      return FilePath

def getDbSize(db):
  command = FastaCmd
  options = ' -I -d ' + db
  p = os.popen(command + options, 'r')
  while 1:
    l = p.readline()
    if not l or 'sequences' in l:
      break
  p.close()
  return int(l.split()[0].replace(',', ''))

def exitOnError(errmsg, exit):
  print 'The following error(s) occurred:\n'
  print errmsg
  sys.exit(exit)

def usage():
  print """Usage:

""",sys.argv[0],""" [-h] -d Database [-t MS Tolerance] [-f] RAW Files MGF FIles

 -h:	Print this help and exit
 -d:	Database containing xlinked peptides
 -t:	Tolerance for MS ions
 -f:	Do not split light and heavy spectra, just search full MSM
 -o:	Output file name (default "Makefile")
 -n:	Experiment name
	"""
  sys.exit(1)

# Paths
SilkConfigPath = '/Users/dawe/Documents/Projects/XL/Silk/trunk/config/'

Raw2XML = which('raw2xml')					# to convert RAW in mzXML files
DoubletCheck = which('doubletCheck.py')		# to split into light and heavy msm
FastaCmd = which('fastacmd')				# to get info about the database
Tolerance = '10'							# default tolerance for doublet checks
Omssa = which('omssacl')
Omssamerge = which ('omssamerge')
Cat = which('cat')
Grep = which('grep')
Rm = which('rm')
Head = which('head')
Awk = which('awk')

# Omssa search parameters
ModificationFile = os.path.join(SilkConfigPath, 'mods.xml')
UserModificationFile = os.path.join(SilkConfigPath, 'usermods.xml')
Missed = '5'
PrecursorTolerance = '0.03'
IonTolerance = '0.5'
HeavyModificationsOpt = ' -mx ' + ModificationFile + ' -mux ' + UserModificationFile + ' -mf 3 -mv 1,120,122'
LightModificationsOpt = ' -mx ' + ModificationFile + ' -mux ' + UserModificationFile + ' -mf 3 -mv 1,119,121'
AllModificationsOpt  = ' -mx ' + ModificationFile + ' -mux ' + UserModificationFile + ' -mf 3 -mv 1,119,121,120,122'
ToleranceOptions = ' -te ' + PrecursorTolerance + ' -to ' + IonTolerance
EnzymeOpts = ' -e 0 -v ' + Missed
MiscOpts = '  -w -zcc 1 -hm 1 -nt 1 -scorp 0.01 -ht 12 -hl 60 -w1 30 -w2 15 -cl 0.01'
ExpectedOpts = ' -he ' #to be modified later on
DbOpts = ' -d ' # to be modified later on

FullSearch = False
OutFile = 'Makefile'
ExperimentName = 'xlink'
MgfExtension = '.msm'	# Default extension for MGF files
RawFiles = []
MgfFiles = []
mzxmlTargets = []
csvTargets = []
lightMgfTargets = []
heavyMgfTargets = []
heavyCsvTargets = []
lightCsvTargets = []

t_all = 'all: '
t_convert = 'convert: '
t_split = 'split: '
t_search = 'search: '
t_light = 'light_search: '
t_heavy = 'heavy_search: '


try:
  optlist, args = getopt.getopt(sys.argv[1:], 'd:t:hfo:n:')
except getopt.GetoptError:
  usage()

if not optlist:
  usage()

for o, a in optlist:
  if o == '-h':
    usage()
  if o == '-d':
    DbIn = a
    DbDir = External.get("blastdb", "path")
    if '/' not in DbIn:
      DbIn = os.path.join(DbDir, DbIn)

  if o == '-t':
    Tolerance = a
  if o == '-f':
    FullSearch = True
  if o == '-o':
    OutFile = a
  if o == '-n':
    ExperimentName = a

t_merge = ExperimentName + ': '
    
if not DbIn:
  exitOnError('Missing database', 1)
else:
  DbOpts += DbIn
  ExpectedOpts += str(100 * getDbSize(DbIn))

# Try to understand which files are RAW and which are MSM
for file in args:
  (fname, ext) = os.path.splitext(file)
  ext = ext.lower()
  if ext == '.raw':
    RawFiles.append(file)
    mzxmlTargets.append(fname + '.mzXML')
  elif ext == '.msm' or ext == '.mgf' or ext == '.msn':
    MgfFiles.append(file)
    lightMgfTargets.append('light_1_' + fname + MgfExtension)
    heavyMgfTargets.append('heavy_1_' + fname + MgfExtension)
    lightCsvTargets.append('light_1_' + fname + '.csv')
    heavyCsvTargets.append('heavy_1_' + fname + '.csv')
    csvTargets.append(fname + '.csv')
    
#print RawFiles
#print mzxmlTargets
#print MgfFiles
#print lightMgfTargets
#print heavyMgfTargets
#print csvTargets

# Time to write the makefile. We have 3 main steps
# 1- convert .RAW files in mzXML (depends on RAW and creates mzXML)
# 2- split .msm files in light and heavy (depends on mzXML and given MSM, gives light and heavy MSM)
# 3- omssa search the light and heavy spectra (depends on light and heavy MSM, gives csv and oms)
# If FullSearch is enabled, there are no step 1 and 2 and step 3 depends on given MSM files
# Additional targets to merge all ommsa search and clean

fh = open(OutFile, 'w') # else is a Makefile

# Start writing all variables
fh.write('RAW2XML = ' + Raw2XML + '\n')
fh.write('DOUBLETCHECK = ' + DoubletCheck + '\n')
fh.write('OMSSA = ' + Omssa + '\n')
fh.write('OMSSAMERGE = ' + Omssamerge + '\n')
fh.write('CAT = ' + Cat + '\n')
fh.write('GREP = ' + Grep + '\n')
fh.write('RM = ' + Rm + '\n')
fh.write('HEAD = ' + Head + '\n')
fh.write('AWK = ' + Awk + '\n')
fh.write('DOUBLETTOL = ' + Tolerance + '\n')
if FullSearch:
  fh.write('MODS = ' + AllModificationsOpt + '\n')
else:
  fh.write('LIGHTMODS = ' + LightModificationsOpt + '\n')
  fh.write('HEAVYMODS = ' + HeavyModificationsOpt + '\n')
fh.write('TOLOPTS = ' + ToleranceOptions + '\n')
fh.write('ENZYMEOPTS = ' + EnzymeOpts + '\n')
fh.write('MISCOPTS = ' + MiscOpts + '\n')
fh.write('EXPECTED = ' + ExpectedOpts + '\n')
fh.write('DB = ' + DbOpts + '\n')
fh.write('MGFEXT = ' + MgfExtension + '\n')
fh.write('\n\n')
# write targets


if FullSearch:
  t_all += 'search ' + ExperimentName
  for csv in csvTargets:
    t_merge = t_merge + ' '  + csv
    t_search = t_search + ' '  + csv

  fh.write(t_all + '\n\ttouch $@\n\n')
  fh.write(t_search  + '\n\ttouch $@\n\n')
  for x, csv in enumerate(csvTargets):
    oms = os.path.splitext(csv)[0] + '.oms'
    fh.write(csv + ': ' + MgfFiles[x] + '\n')
    fh.write('\t$(OMSSA) $(MODS) $(TOLOPTS) $(ENZYMEOPTS) $(MISCOPTS) $(EXPECTED) $(DB) -oc $@ -fm $< -ob ' + oms + '\n\n')
else:
  t_all += 'convert split search ' + ExperimentName
  for xml in mzxmlTargets:
    t_convert = t_convert + ' ' + xml
  for file in lightMgfTargets + heavyMgfTargets:
    t_split = t_split + ' ' + file

  t_search += 'light_search heavy_search'
  for csv in heavyCsvTargets: 
    t_heavy = t_heavy + ' ' + csv
  for csv in lightCsvTargets:
    t_light = t_light + ' ' + csv
  
  fh.write(t_all + '\n\ttouch $@\n\n')
  fh.write(t_convert + '\n\ttouch $@\n\n')
  fh.write(t_split + '\n\ttouch $@\n\n')
  fh.write(t_search  + '\n\ttouch $@\n\n')
  fh.write(t_light   + '\n\ttouch $@\n\n')
  fh.write(t_heavy + '\n\ttouch $@\n\n')
  
  
  # Conversion
  for x, xml in enumerate(mzxmlTargets):
    fh.write(xml + ': ' + RawFiles[x] + '\n')
    fh.write('\t $(RAW2XML) --mzXML $< $@\n\n')
  
  # Split
  for x, msm in enumerate(lightMgfTargets):
    fh.write(msm + ': ' + mzxmlTargets[x] + '\n')
    fh.write('\t$(DOUBLETCHECK) -3 -t ' + Tolerance + ' -f ' + MgfFiles[x] + ' -m ' + mzxmlTargets[x] + '\n\n')
    fh.write(heavyMgfTargets[x] + ': ' + msm + '\n')
    fh.write('\ttouch $@\n\n')
  # OMSSA
  for x, csv in enumerate(lightCsvTargets):
    t_merge = t_merge + ' '  + csv
    oms = os.path.splitext(csv)[0] + '.oms'
    fh.write(csv + ': ' + lightMgfTargets[x] + '\n')
    fh.write('\t$(OMSSA) $(LIGHTMODS) $(TOLOPTS) $(ENZYMEOPTS) $(MISCOPTS) $(EXPECTED) $(DB) -oc $@ -fm $< -ob ' + oms + '\n\n')
  for x, csv in enumerate(heavyCsvTargets):
    t_merge = t_merge + ' '  + csv
    oms = os.path.splitext(csv)[0] + '.oms'
    fh.write(csv + ': ' + heavyMgfTargets[x] + '\n')
    fh.write('\t$(OMSSA) $(HEAVYMODS) $(TOLOPTS) $(ENZYMEOPTS) $(MISCOPTS) $(EXPECTED) $(DB) -oc $@ -fm $< -ob ' + oms + '\n\n')
    
# Merge results
fh.write(t_merge + '\n')
Commands = """	@echo "Merging csv files... "
	$(HEAD) -n1 $< > $@.csv
	for i in $?; do \
		mfile=$${i%%.csv}$(MGFEXT); \
		thisn=`grep BEGIN $$mfile|wc -l`; \
		$(GREP) -v 'Spectrum number,' $$i | $(AWK) -v add=$$n 'BEGIN {FS=","; OFS=","}; {$$1=$$1 + add; print $$0}' >>$@.csv ; \
		n=$$(($$n + $$thisn)) ; \
	done;
	@echo "Merging $(MGFEXT) files... "
	$(HEAD) -n2 $(subst .csv,$(MGFEXT),$<) > $@$(MGFEXT)
	$(GREP) -h -v ^[SR] $(subst .csv,$(MGFEXT),$?) >> $@$(MGFEXT)
	@echo "Merging oms files..."
	ls -1 $(subst .csv,.oms,$?) > $@.list
	$(OMSSAMERGE) -i $@.list -ib -ob $@.oms
	$(RM) -f $@.list

"""
fh.write(Commands)
fh.write('\ttouch ' + ExperimentName + '\n\n')

# Clean
Commands = """.PHONY: clean
clean:
	$(RM) -f *.csv *.oms test.dta """ + ExperimentName + """* all search light_search heavy_search convert split
	
deep_clean:
	$(RM) -f *.mzXML light* heavy* *.csv *.oms test.dta """ + ExperimentName + """* all search light_search heavy_search convert split
"""
fh.write(Commands)
